# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""异步函数测试

测试async/await语法的各种场景。
"""

from gtest import EXPECT_EQ, EXPECT_TRUE, EXPECT_FALSE, test

from mc.future import sleep


@test("Async", "BasicAsyncFunction")
async def test_basic_async_function():
    """测试基础async函数定义（无await）"""

    async def get_value() -> int:
        return 42

    EXPECT_EQ(await get_value(), 42)


# 全局辅助函数（避免闭包捕获问题）
async def async_step1() -> int:
    return 10


async def async_step2(x: int) -> int:
    y = await async_step1()
    return x * 2 + y


@test("Async", "ThenChainTest")
async def test_then_chain():
    """测试Then链串联转换"""

    a = await async_step1()
    b = await async_step2(a)

    # async_step1() = 10, async_step2(10) = 10*2 + 10 = 30
    EXPECT_EQ(b, 30)


@test("Async", "NestedAsyncFunctions")
async def test_nested_async_functions():
    """测试嵌套async函数相互调用"""

    async def inner1() -> int:
        return 10

    async def inner2(x: int) -> int:
        return x * 2

    a = await inner1()
    b = await inner2(a)
    c = await async_step2(b)

    # inner1() = 10, inner2(10) = 20, async_step2(20) = 20*2 + 10 = 50
    EXPECT_EQ(c, 50)


# 测试单个 await
@test("StacklessBasic", "SingleAwait")
async def test_single_await():
    """测试单个 await 的情况"""

    async def step() -> int:
        return 10

    # 测试直接调用（有明确类型）
    async def compute(count: int) -> int:
        a = 0
        for _ in range(count):
            a += 1
        v = await step()
        return v + 5 + a

    EXPECT_EQ(await compute(100), 115)


@test("Async", "CodeBetweenAwaits")
async def test_code_between_awaits():
    """测试await之间插入普通代码"""

    async def step1() -> int:
        return 5

    async def step2(x: int) -> int:
        return x + 3

    async def step3(x: int) -> int:
        return x * 2

    # 第一个await
    a = await step1()
    temp = a + 1

    # 第二个await
    b = await step2(temp)
    result = b * 3

    # 第三个await
    c = await step3(result)

    final = c + 10

    EXPECT_EQ(final, 64)


@test("Async", "DeepRecursion")
async def test_deep_recursion():
    """测试超过32个连续同步await的递归深度限制"""

    # add_one 函数立即返回并没有真正的异步操作，运行时引擎会优化为同步递归调用
    async def add_one(x: int) -> int:
        return x + 1

    # 创建40个连续的await调用（超过max_sync_recursion_depth=32）
    v = 0
    v = await add_one(v)  # 1
    v = await add_one(v)  # 2
    v = await add_one(v)  # 3
    v = await add_one(v)  # 4
    v = await add_one(v)  # 5
    v = await add_one(v)  # 6
    v = await add_one(v)  # 7
    v = await add_one(v)  # 8
    v = await add_one(v)  # 9
    v = await add_one(v)  # 10
    v = await add_one(v)  # 11
    v = await add_one(v)  # 12
    v = await add_one(v)  # 13
    v = await add_one(v)  # 14
    v = await add_one(v)  # 15
    v = await add_one(v)  # 16
    v = await add_one(v)  # 17
    v = await add_one(v)  # 18
    v = await add_one(v)  # 19
    v = await add_one(v)  # 20
    v = await add_one(v)  # 21
    v = await add_one(v)  # 22
    v = await add_one(v)  # 23
    v = await add_one(v)  # 24
    v = await add_one(v)  # 25
    v = await add_one(v)  # 26
    v = await add_one(v)  # 27
    v = await add_one(v)  # 28
    v = await add_one(v)  # 29
    v = await add_one(v)  # 30
    v = await add_one(v)  # 31
    v = await add_one(v)  # 32
    v = await add_one(v)  # 33 - 超过最大递归阈值，应该触发异步化
    v = await add_one(v)  # 34
    v = await add_one(v)  # 35
    v = await add_one(v)  # 36
    v = await add_one(v)  # 37
    v = await add_one(v)  # 38
    v = await add_one(v)  # 39
    v = await add_one(v)  # 40

    # 期望结果：0 + 40 = 40
    EXPECT_EQ(v, 40)


# ============================================================================
# 控制流测试
# ============================================================================


@test("AsyncControlFlow", "IfWithAwaitInThen")
async def test_if_with_await_in_then():
    """测试 if 分支中的 await"""

    async def get_value(x: int) -> int:
        return x

    condition = True
    result = 0

    if condition:
        result = await get_value(10)

    EXPECT_EQ(result, 10)


@test("AsyncControlFlow", "IfWithAwaitInElse")
async def test_if_with_await_in_else():
    """测试 else 分支中的 await"""

    async def get_value(x: int) -> int:
        return x

    condition = False
    result = 0

    if condition:
        result = 5
    else:
        result = await get_value(20)

    EXPECT_EQ(result, 20)


@test("AsyncControlFlow", "IfWithAwaitInBoth")
async def test_if_with_await_in_both():
    """测试 if 和 else 分支都有 await"""

    async def get_value(x: int) -> int:
        return x

    async def process(use_then: bool) -> int:
        if use_then:
            x = await get_value(10)
            return x
        else:
            y = await get_value(20)
            return y

    EXPECT_EQ(await process(True), 10)
    EXPECT_EQ(await process(False), 20)


@test("AsyncControlFlow", "ForLoopWithAwait")
async def test_for_loop_with_await():
    """测试 for 循环中的 await"""

    async def get_value(x: int) -> int:
        val = await async_step1()
        return x + val

    total = 0
    for i in range(3):
        val = await get_value(i)
        total = total + val

    # (0+10) + (1+10) + (2+10) = 33
    EXPECT_EQ(total, 33)


@test("AsyncControlFlow", "WhileLoopWithAwait")
async def test_while_loop_with_await():
    """测试 while 循环中的 await"""

    async def get_value(x: int) -> int:
        val = await async_step1()
        return x + val

    total = 0
    i = 0
    while i < 3:
        val = await get_value(i)
        total = total + val
        i = i + 1

    # (0+10) + (1+10) + (2+10) = 33
    EXPECT_EQ(total, 33)


@test("AsyncControlFlow", "NestedControlFlowWithAwait")
async def test_nested_control_flow_with_await():
    """测试嵌套控制流中的 await"""

    async def get_value(x: int) -> int:
        base = await async_step1()
        return x + base

    total = 0
    for i in range(2):
        if i == 0:
            val = await get_value(5)
            total = total + val
        else:
            val = await get_value(10)
            total = total + val

    # (5+10) + (10+10) = 35
    EXPECT_EQ(total, 35)


@test("AsyncControlFlow", "AwaitBeforeIfWithAwaitInBranches")
async def test_await_before_if_with_await_in_branches():
    """测试 if 之前有 await，分支中也有 await（会导致 cond 变量提升）"""

    async def get_true() -> bool:
        return True

    async def get_value_a() -> int:
        return 10

    async def get_value_b() -> int:
        return 20

    # 这个模式会导致 cond 变量需要提升为状态变量
    # 因为 cond 赋值之后，分支中有 await
    result: int = 0
    cond = await get_true()

    if cond:
        result = await get_value_a()
    else:
        result = await get_value_b()

    EXPECT_EQ(result, 10)


@test("AsyncControlFlow", "ForLoopWithBreak")
async def test_for_loop_with_break():
    """测试 for 循环中的 break"""

    async def get_value(x: int) -> int:
        val = await async_step1()
        return x + val

    total = 0
    for i in range(10):
        val = await get_value(i)
        if val > 15:
            break
        total = total + val

    # (0+10) + (1+10) + (2+10) + (3+10) + (4+10) + (5+10) = 10+11+12+13+14+15 = 75
    # 当 i=6 时，val=16 > 15，break
    EXPECT_EQ(total, 75)


@test("AsyncControlFlow", "ForLoopWithContinue")
async def test_for_loop_with_continue():
    """测试 for 循环中的 continue"""

    async def get_value(x: int) -> int:
        val = await async_step1()
        return x + val

    total = 0
    for i in range(5):
        val = await get_value(i)
        if val % 2 == 0:
            continue
        total = total + val

    # i=0: val=10(偶数,跳过), i=1: val=11(奇数,累加), i=2: val=12(偶数,跳过)
    # i=3: val=13(奇数,累加), i=4: val=14(偶数,跳过)
    # total = 11 + 13 = 24
    EXPECT_EQ(total, 24)


@test("AsyncControlFlow", "WhileLoopWithBreak")
async def test_while_loop_with_break():
    """测试 while 循环中的 break"""

    async def get_value(x: int) -> int:
        val = await async_step1()
        return x + val

    total = 0
    i = 0
    while i < 10:
        val = await get_value(i)
        if val > 15:
            break
        total = total + val
        i = i + 1

    # (0+10) + (1+10) + (2+10) + (3+10) + (4+10) + (5+10) = 10+11+12+13+14+15 = 75
    # 当 i=6 时，val=16 > 15，break
    EXPECT_EQ(total, 75)


@test("AsyncControlFlow", "WhileLoopWithContinue")
async def test_while_loop_with_continue():
    """测试 while 循环中的 continue"""

    async def get_value(x: int) -> int:
        val = await async_step1()
        return x + val

    total = 0
    i = 0
    while i < 5:
        val = await get_value(i)
        i = i + 1
        if val % 2 == 0:
            continue
        total = total + val

    # i=0: val=10(偶数,跳过), i=1: val=11(奇数,累加), i=2: val=12(偶数,跳过)
    # i=3: val=13(奇数,累加), i=4: val=14(偶数,跳过)
    # total = 11 + 13 = 24
    EXPECT_EQ(total, 24)


@test("AsyncControlFlow", "NestedLoopWithBreak")
async def test_nested_loop_with_break():
    """测试嵌套循环中的 break（只跳出内层循环）"""

    async def get_value(x: int, y: int) -> int:
        val = await async_step1()
        return x + y + val

    total = 0
    for i in range(3):
        for j in range(5):
            val = await get_value(i, j)
            if j >= 2:
                break
            total = total + val

    # i=0: j=0 (val=10), j=1 (val=11), j=2 break
    # i=1: j=0 (val=11), j=1 (val=12), j=2 break
    # i=2: j=0 (val=12), j=1 (val=13), j=2 break
    # total = 10 + 11 + 11 + 12 + 12 + 13 = 69
    EXPECT_EQ(total, 69)


# # ============================================================================
# # await 在表达式中的测试
# # ============================================================================


@test("AwaitInExpression", "BinaryOperation")
async def test_await_in_binary_operation():
    """测试 await 在二元运算中的使用"""

    async def get_value(x: int) -> int:
        return x

    # await 在二元运算中
    result = await get_value(5) + 1

    EXPECT_EQ(result, 6)


@test("AwaitInExpression", "FunctionArgument")
async def test_await_in_function_argument():
    """测试 await 作为函数参数"""

    async def get_value(x: int) -> int:
        return x

    async def add(a: int, b: int) -> int:
        return a + b

    # await 作为函数参数
    result = await add(await get_value(3), await get_value(4))

    EXPECT_EQ(result, 7)


@test("AwaitInExpression", "NestedAwait")
async def test_nested_await():
    """测试嵌套的 await"""

    async def get_value(x: int) -> int:
        return x

    async def double_value(x: int) -> int:
        return x * 2

    # 嵌套 await
    result = await double_value(await get_value(5))

    EXPECT_EQ(result, 10)


@test("AwaitInExpression", "MultipleOperations")
async def test_await_multiple_operations():
    """测试多个 await 在复杂表达式中"""

    async def get_value(x: int) -> int:
        return x

    # 多个 await 在一个表达式中
    result = await get_value(1) + await get_value(2) + await get_value(3)

    EXPECT_EQ(result, 6)


# ============================================================================
# 短路求值测试
# ============================================================================
# 注意：三元表达式的短路求值暂时有作用域问题，暂时跳过


@test("ShortCircuit", "AndBothFalse")
async def test_and_both_false():
    """测试 and：两个操作数都为 False"""

    async def get_false1() -> bool:
        return False

    async def get_false2() -> bool:
        return False

    result = await get_false1() and await get_false2()
    EXPECT_FALSE(result)


@test("ShortCircuit", "AndTrueAndFalse")
async def test_and_true_and_false():
    """测试 and：左True右False"""

    async def get_true() -> bool:
        return True

    async def get_false() -> bool:
        return False

    result = await get_true() and await get_false()
    EXPECT_FALSE(result)


@test("ShortCircuit", "AndBothTrue")
async def test_and_both_true():
    """测试 and：两个操作数都为 True"""

    async def get_true1() -> bool:
        return True

    async def get_true2() -> bool:
        return True

    result = await get_true1() and await get_true2()
    EXPECT_TRUE(result)


@test("ShortCircuit", "OrBothFalse")
async def test_or_both_false():
    """测试 or：两个操作数都为 False"""

    async def get_false1() -> bool:
        return False

    async def get_false2() -> bool:
        return False

    result = await get_false1() or await get_false2()
    EXPECT_FALSE(result)


@test("ShortCircuit", "OrTrueAndFalse")
async def test_or_true_and_false():
    """测试 or：左True右False"""

    async def get_true() -> bool:
        return True

    async def get_false() -> bool:
        return False

    result = await get_true() or await get_false()

    # 手写短路求值对比
    result1 = await get_true()
    if not result1:
        result1 = await get_false()

    EXPECT_TRUE(result == result1)


@test("ShortCircuit", "OrBothTrue")
async def test_or_both_true():
    """测试 or：两个操作数都为 True"""

    async def get_true1() -> bool:
        return True

    async def get_true2() -> bool:
        return True

    result = await get_true1() or await get_true2()
    EXPECT_TRUE(result)


@test("ShortCircuit", "ConditionalTrueBranch")
async def test_conditional_true_branch():
    """测试条件表达式：条件为 True"""

    async def get_true() -> bool:
        return True

    async def get_value_a() -> int:
        return 10

    async def get_value_b() -> int:
        return 20

    result = await get_value_a() if await get_true() else await get_value_b()
    EXPECT_EQ(result, 10)


@test("ShortCircuit", "ConditionalFalseBranch")
async def test_conditional_false_branch():
    """测试条件表达式：条件为 False"""

    async def get_false() -> bool:
        return False

    async def get_value_a() -> int:
        return 10

    async def get_value_b() -> int:
        return 20

    result = await get_value_a() if await get_false() else await get_value_b()
    EXPECT_EQ(result, 20)


@test("ShortCircuit", "ConditionalMixedExpression")
async def test_conditional_mixed_expression():
    """测试三元表达式：条件、true分支、false分支都包含混合表达式（异步+同步+运算符）"""

    async def async_get_value(x: int) -> int:
        return x

    def sync_double(x: int) -> int:
        return x * 2

    async def async_check_positive(x: int) -> bool:
        return x > 0

    def sync_add(a: int, b: int) -> int:
        return a + b

    async def compute(check_value: int) -> int:
        # 复杂的三元表达式：
        # - 条件：await async_check_positive(await async_get_value(check_value)) and sync_add(1, 1) == 2
        # - true分支：await async_get_value(10) + sync_double(3)
        # - false分支：await async_get_value(20) * 2 + sync_add(5, 5)
        result = (
            await async_get_value(10) + sync_double(3)
            if await async_check_positive(await async_get_value(check_value))
            and sync_add(1, 1) == 2
            else await async_get_value(20) * 2 + sync_add(5, 5)
        )
        return result

    # 测试条件为 True 的情况：check_value=5 > 0
    # 执行 true 分支：async_get_value(10) + sync_double(3) = 10 + 6 = 16
    result_true = await compute(5)
    EXPECT_EQ(result_true, 16)

    # 测试条件为 False 的情况：check_value=-5 < 0
    # 执行 false 分支：async_get_value(20) * 2 + sync_add(5, 5) = 20 * 2 + 10 = 50
    result_false = await compute(-5)
    EXPECT_EQ(result_false, 50)


# ============================================================================
# 混合普通控制流和异步数据流测试
# ============================================================================


@test("MixedControlFlow", "ForLoopWithoutAwaitBeforeAwait")
async def test_for_loop_without_await_before_await():
    """测试不包含 await 的 for 循环在 await 之前"""

    async def step() -> int:
        return 10

    a = 0
    for _ in range(100):
        a += 1
    v = await step()

    EXPECT_EQ(v + 5 + a, 115)


@test("MixedControlFlow", "IfWithoutAwaitBeforeAwait")
async def test_if_without_await_before_await():
    """测试不包含 await 的 if 语句在 await 之前"""

    async def get_value() -> int:
        return 20

    async def process(condition: bool) -> int:
        result = 0
        if condition:
            result = 5
        else:
            result = 10
        v = await get_value()
        return v + result

    EXPECT_EQ(await process(True), 25)
    EXPECT_EQ(await process(False), 30)


@test("MixedControlFlow", "NestedLoopsWithoutAwait")
async def test_nested_loops_without_await():
    """测试嵌套的不包含 await 的循环"""

    async def get_value() -> int:
        return 100

    total = 0
    for i in range(3):
        for j in range(2):
            total = total + i + j
    v = await get_value()

    # i=0: j=0(0+0=0), j=1(0+1=1) -> 1
    # i=1: j=0(1+0=1), j=1(1+1=2) -> 3
    # i=2: j=0(2+0=2), j=1(2+1=3) -> 5
    # total = 1 + 3 + 5 = 9
    # result = 100 + 9 = 109
    EXPECT_EQ(v + total, 109)


@test("MixedControlFlow", "MixedControlFlowWithAwait")
async def test_mixed_control_flow_with_await():
    """测试混合的控制流：部分包含 await，部分不包含"""

    async def get_value(x: int) -> int:
        return x + 10

    total = 0
    # 不包含 await 的循环
    for i in range(5):
        total = total + i

    # await 语句
    v = await get_value(total)

    # 不包含 await 的 if 语句
    result: int = 0
    if v > 50:
        result = v * 2
    else:
        result = v

    # total = 0+1+2+3+4 = 10, v = 10+10 = 20, result = 20 (因为 20 <= 50)
    EXPECT_EQ(result, 20)


@test("MixedControlFlow", "AwaitBetweenControlFlows")
async def test_await_between_control_flows():
    """测试 await 在控制流语句之间"""

    async def get_value(x: int) -> int:
        return x + 5

    # 第一个不包含 await 的循环
    sum1 = 0
    for i in range(3):
        sum1 = sum1 + i

    # await 语句
    v = await get_value(sum1)

    # 第二个不包含 await 的循环
    sum2 = 0
    for j in range(2):
        sum2 = sum2 + j

    # 不包含 await 的 if 语句
    result: int = 0
    if v > 10:
        result = v + sum2
    else:
        result = v - sum2

    # sum1 = 0+1+2 = 3, v = 3+5 = 8, sum2 = 0+1 = 1
    # v = 8 <= 10, result = 8 - 1 = 7
    EXPECT_EQ(result, 7)


@test("MixedControlFlow", "ComplexMixedControlFlow")
async def test_complex_mixed_control_flow():
    """测试复杂的混合控制流场景（包含变量跨 await 使用）"""

    async def async_step(x: int) -> int:
        return x * 2

    # 初始化（第一个 Unit）
    a = 0
    b = 0

    # 不包含 await 的循环
    for i in range(3):
        a = a + i

    # await 语句
    v1 = await async_step(a)

    # 不包含 await 的 if 语句（第二个 Unit）
    if v1 > 5:
        b = 10
    else:
        b = 20

    # 不包含 await 的嵌套循环
    c = 0
    for i in range(2):
        for j in range(2):
            c = c + i + j

    # 第二个 await 语句
    v2 = await async_step(c)

    # 不包含 await 的 if 语句（第三个 Unit）
    result: int = 0
    if v2 > 10:
        result = v1 + v2 + b
    else:
        result = v1 + v2 - b

    # a = 0+1+2 = 3, v1 = 3*2 = 6
    # v1 = 6 > 5, b = 10
    # c = (0+0) + (0+1) + (1+0) + (1+1) = 0+1+1+2 = 4
    # v2 = 4*2 = 8
    # v2 = 8 <= 10, result = 6 + 8 - 10 = 4
    EXPECT_EQ(result, 4)


# ============================================================================
# 测试异步函数延迟执行
# ============================================================================


@test("Async", "DelayedExecution")
async def test_delayed_execution():
    """测试异步函数延迟执行"""

    async def sleep_ms(duration_ms: int):
        await sleep(duration_ms)

    await sleep(10)
    await sleep_ms(10)
    result = 10

    EXPECT_EQ(result, 10)
